import numpy as np
import operator
import pickle

from treePlotter import retrieveTree, createPlot

# 计算给定数据集的香农熵,输入数据不能使用numpy数组，因为含有字符串
def calcShannonEnt(dataSet):
    numEntries = len(dataSet)
    labelCounts = {}
    for featVec in dataSet:
        currentLabel = featVec[-1]
        labelCounts[currentLabel]=labelCounts.get(currentLabel,0)+1
    # print(labelCounts)
    shannonEnt = 0
    for key in labelCounts:
        prob = labelCounts[key]/float(numEntries)
        shannonEnt += -prob*np.log2(prob) 
    return shannonEnt

# 按照给定特征划分数据集
def splitDataSet(dataSet, axis, value):
    retDataSet = []
    for featVec in dataSet:
        if featVec[axis]==value:
            reducedFeatVec = featVec[:axis]
            reducedFeatVec.extend(featVec[axis+1:])
            retDataSet.append(reducedFeatVec)
    return retDataSet

# 选择最好的数据集划分方式
def chooseBestFeatureToSplit(dataSet):
    numFeatures = len(dataSet[0])-1
    baseEntropy = calcShannonEnt(dataSet)
    bestInfoGain=0
    bestFeature = -1
    for i in range(numFeatures):
        featList = [example[i] for example in dataSet]
        uniqueVals = set(featList)
        newEntropy = 0
        for value in uniqueVals:
            subDataSet = splitDataSet(dataSet,i,value)
            prob = len(subDataSet)/float(len(dataSet))
            newEntropy += prob* calcShannonEnt(subDataSet)
        infoGain = baseEntropy-newEntropy
        if infoGain> bestInfoGain:
            bestInfoGain = infoGain
            bestFeature = i
    return bestFeature

# 采用多数表决的方法，决定一个节点的分类
def majorityCnt(classList):
    classCount = {}
    for vote in classList:
        classCount[vote]=classCount.get(vote,0) + 1
    
    # sortedClassCount = sorted(classCount.items(),operator.itemgetter(1),reverse=True)
    temp = sorted(classCount.items(),key=operator.itemgetter(1),reverse=True)
    return temp[0][0]

def createTree(dataSet,labels):
    classList = [example[-1] for example in dataSet]
    # 表明类别完全相同
    if classList.count(classList[0]) == len(classList):
        return classList[0]
    # 使用完了所有特征，仍然包含类别不唯一，返回类别较多的那个
    if len(dataSet[0])==1:
        return majorityCnt(classList)
    
    bestFeat = chooseBestFeatureToSplit(dataSet)
    bestFeatLabel = labels[bestFeat]
    
    myTree = {bestFeatLabel:{}}
    del(labels[bestFeat])
    
    featValues = [example[bestFeat] for example in dataSet]
    uniqueVals = set(featValues) # 0 1
    
    for value in uniqueVals:
        subLabels = labels[:]
        subDataSet = splitDataSet(dataSet,bestFeat,value)
        
        myTree[bestFeatLabel][value]=createTree(subDataSet,subLabels)
    return myTree

# 使用决策树的分类器
def classify(inputTree, featLabels, testVec):
    firstStr = list(inputTree.keys())[0]
    secondDict = inputTree[firstStr]

    featIndex = featLabels.index(firstStr)

    for key in secondDict.keys():
        if testVec[featIndex] == key:
            if type(secondDict[key]).__name__ == 'dict':
                classLabel = classify(secondDict[key],featLabels,testVec)
            else:
                classLabel = secondDict[key]
    return classLabel

# 决策树的存储
def storeTree(inputTree,filename):
    fw = open(filename, 'wb')
    pickle.dump(inputTree,fw)
    fw.close()

def grabTree(filename):
    fr = open(filename,'rb')
    return pickle.load(fr)


def createDataSet():
    dataSet = [
        [1,1,'yes'],
        [1,1,'yes'],
        [1,0,'no'],
        [0,1,'no'],
        [0,1,'no']
    ]
    labels = ['no surfacing', 'flippers']
    return dataSet,labels
if __name__ == '__main__':
    
    # myDat, labels = createDataSet()
    # testVec=[1,1]
    # tree = grabTree('ts.txt')
    # print(tree)
    # ret = classify(tree, labels, testVec) 
    # print(ret)

    fr = open('data/lenses.txt')
    lenses = [inst.strip().split('\t') for inst in fr.readlines()]
    lensesLabels =['age', 'prescript', 'astigmatic', 'tearRate']

    lensesTree = createTree(lenses, lensesLabels)
    print(lensesTree)
    createPlot(lensesTree)
    
    

    



    